# coding: utf-8
import os
DATASET_ROOT_DIR = "~/dataset/voc_road"
DATASET_ROOT_DIR = os.path.expanduser(DATASET_ROOT_DIR)


result_lable = {
    0: "bk",
    1: "person",
    2: "car",
    3: "bicycle",
    4: "motorbike"
}

voc_label = {
    "person": 1,
    "bicycle": 3,
    "motorbike": 4,
    "car": 2
}

coco_label = {
    1: "person",
    2:"bicycle",
    3: "car",
    4: "motorbike"
}

def mk_common_dir(root_path):
    ## Annotations
    fp = os.path.join(root_path, "Annotations")
    if not os.path.exists(fp):
        os.mkdir(fp)

    ## ImageSets
    fp = os.path.join(root_path, "ImageSets")
    if not os.path.exists(fp):
        os.mkdir(fp)
    ## ImageSets/Main
    fp = fp + "/Main"
    if not os.path.exists(fp):
        os.mkdir(fp)

    ## JPEGImages
    fp = os.path.join(root_path, "JPEGImages")
    if not os.path.exists(fp):
        os.mkdir(fp)

from xml.dom.minidom import Document
def write_xml(xml_path, image_name, folder_source, image_size, bboxes, object_names):
    xmlpath = xml_path + "/" + image_name + ".xml"
    doc = Document()
    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)
    folder = doc.createElement('folder')
    folder_name = doc.createTextNode(folder_source)
    folder.appendChild(folder_name)
    annotation.appendChild(folder)
    filenamenode = doc.createElement('filename')
    filename_name = doc.createTextNode(image_name + ".jpg")
    filenamenode.appendChild(filename_name)
    annotation.appendChild(filenamenode)
    # source
    source = doc.createElement('source')
    annotation.appendChild(source)
    database = doc.createElement('database')
    database.appendChild(doc.createTextNode(folder_source))
    source.appendChild(database)
    annotation_s = doc.createElement('annotation')
    annotation_s.appendChild(doc.createTextNode(folder_source))
    source.appendChild(annotation_s)
    image = doc.createElement('image')
    image.appendChild(doc.createTextNode(folder_source))
    source.appendChild(image)

    # flickrid = doc.createElement('flickrid')
    # flickrid.appendChild(doc.createTextNode('-1'))
    # source.appendChild(flickrid)
    # owner = doc.createElement('owner')
    # annotation.appendChild(owner)
    # flickrid_o = doc.createElement('flickrid')
    # flickrid_o.appendChild(doc.createTextNode('yanyu'))
    # owner.appendChild(flickrid_o)
    # name_o = doc.createElement('name')
    # name_o.appendChild(doc.createTextNode('yanyu'))
    # owner.appendChild(name_o)

    # size
    size = doc.createElement('size')
    annotation.appendChild(size)
    width = doc.createElement('width')
    width.appendChild(doc.createTextNode(str(image_size[0])))
    height = doc.createElement('height')
    height.appendChild(doc.createTextNode(str(image_size[1])))
    depth = doc.createElement('depth')
    depth.appendChild(doc.createTextNode(str(image_size[2])))
    size.appendChild(width)
    size.appendChild(height)
    size.appendChild(depth)

    # segment
    segmented = doc.createElement('segmented')
    segmented.appendChild(doc.createTextNode('0'))
    annotation.appendChild(segmented)

    # object
    for i in range(len(bboxes)):
        bbox = bboxes[i]
        name = object_names[i]
        objects = doc.createElement('object')
        annotation.appendChild(objects)
        object_name = doc.createElement('name')
        object_name.appendChild(doc.createTextNode(name))
        objects.appendChild(object_name)
        pose = doc.createElement('pose')
        pose.appendChild(doc.createTextNode('Unspecified'))
        objects.appendChild(pose)
        truncated = doc.createElement('truncated')
        truncated.appendChild(doc.createTextNode('1'))
        objects.appendChild(truncated)
        difficult = doc.createElement('difficult')
        difficult.appendChild(doc.createTextNode('0'))
        objects.appendChild(difficult)
        bndbox = doc.createElement('bndbox')
        objects.appendChild(bndbox)
        xmin = doc.createElement('xmin')
        xmin.appendChild(doc.createTextNode(str(bbox[0])))
        bndbox.appendChild(xmin)
        ymin = doc.createElement('ymin')
        ymin.appendChild(doc.createTextNode(str(bbox[1])))
        bndbox.appendChild(ymin)
        xmax = doc.createElement('xmax')
        xmax.appendChild(doc.createTextNode(str(bbox[2])))
        bndbox.appendChild(xmax)
        ymax = doc.createElement('ymax')
        ymax.appendChild(doc.createTextNode(str(bbox[3])))
        bndbox.appendChild(ymax)
    f = open(xmlpath, "w")
    f.write(doc.toprettyxml(indent=''))
    f.close()
